package com.tencent.angel.sona.ml.common

/**
  * @author xiongjun
  * @date 2019/12/13 15:19
  * @description
  * @reviewer
  */
import org.apache.spark.TaskContext
import MathImplicits._
import com.tencent.angel.sona.ml.math2.utils.CusLabeledData
import org.apache.spark.util.Example
import org.apache.spark.rdd.RDD
import org.apache.spark.storage.StorageLevel

import scala.reflect.ClassTag


/**
  * Build manifold view for a RDD. A manifold RDD is to split a RDD to multiple RDD.
  * First, we shuffle the RDD and split it into several splits inside every partition.
  * Then, we hold the manifold RDD into cache.
  */

class CusManifoldBuilder(data: RDD[Example],
                      numSplit: Int,
                      partitionStat: Map[Int, Long],
                      persist: StorageLevel = StorageLevel.MEMORY_AND_DISK,
                      preservesPartitioning: Boolean = true)(implicit dim: Long) extends Serializable {
  protected def trans(item: Example): CusLabeledData = {
    new CusLabeledData(item.features, item.label)
  }

  private def split(iterator: Iterator[Example]): Iterator[Array[CusLabeledData]] = {
    new ManifoldSplitter(iterator, numSplit, trans, partitionStat)
  }

  lazy val foldedRDD: RDD[Array[CusLabeledData]] = data.mapPartitions(
    itr => split(itr), preservesPartitioning)

  def manifoldRDD(): Seq[RDD[Array[CusLabeledData]]] = {
    Vector.tabulate(numSplit) { i =>
      foldedRDD.mapPartitions(itr => Iterator.single(itr.drop(i).next()), preservesPartitioning)
    }
  }
}

